#include "resource/resource.hpp"
#include <algorithm>
#include <cstddef>
#include <cstdlib>
#include <iterator>
#include <string>
#include <unordered_map>
#include <vector>
#include "core/assert.hpp"
#include "core/log.hpp"
#include "utils/string_utils.hpp"

#include "zstd.h"

enum fs_type {
    file,
    dir,
};

struct fs_node_base {
    fs_type     type; // 如果是dir，则可以有一些子节点
};

struct fs_dir : fs_node_base {
    std::unordered_map<std::string, fs_node_base*> children;
};

struct fs_file: fs_node_base {
    const void *data_ptr;
    size_t      data_length;
    size_t      origin_length;
};

struct embedded_resource_data
{
    // root node
    fs_dir *root = nullptr;

    embedded_resource_data() {
        root = new fs_dir;
        root->type = fs_type::dir;
    }

    fs_dir *get_dir(const std::string &dir) const {
        auto paths = parse_path(dir);
        fs_dir *result = root;
        for (auto it = begin(paths); it != end(paths); ++it) {
            // recurse find dir
            auto it_child = result->children.find(*it);
            if (it_child != end(result->children) && it_child->second->type == fs_type::dir) {
                result = static_cast<fs_dir*>(it_child->second);
            } else {
                // not found
                result = nullptr;
                break;
            }
        }
        return result;
    }

    fs_file *get_file(const std::string &path) const {
        auto paths = parse_path(path);
        fs_dir  *dir = root;
        fs_file *result = nullptr;
        // check dir exists
        for (size_t i = 0; i<paths.size()-1; ++i) {
            auto it_child = dir->children.find(paths[i]);
            if (it_child != end(dir->children) && it_child->second->type == fs_type::dir) {
                dir = static_cast<fs_dir*>(it_child->second);
            } else {
                // not found
                dir = nullptr;
                break;
            }
        }
        if (dir != nullptr) {
            // check file exists
            auto it_child = dir->children.find(paths.back());
            if (it_child != end(dir->children) && it_child->second->type == fs_type::file) {
                result = static_cast<fs_file*>(it_child->second);
            }
        }
        return result;
    }

    bool set_file(const std::string &path, const void *data, size_t length, size_t origin_length)
    {
        auto paths = parse_path(path);
        fs_dir  *dir = root;
        for (size_t i = 0; i<paths.size()-1; ++i) {
            auto it_child = dir->children.find(paths[i]);
            if (it_child != end(dir->children) && it_child->second->type == fs_type::dir) {
                dir = static_cast<fs_dir*>(it_child->second);
            } else if (it_child == end(dir->children)) {
                // not found, create it
                auto const new_dir = new fs_dir;
                new_dir->type = fs_type::dir;
                auto const ret = dir->children.emplace(std::pair{paths[i], static_cast<fs_node_base*>(new_dir)});
                PI_ASSERT(ret.second);
                dir = static_cast<fs_dir*>(ret.first->second);
            } else {
                L_ERROR("resource %s(%u) exists, but not a directory!!", path.c_str(), i);
                return false;
            }
        }
        // check file not exists
        if (auto it = dir->children.find(paths.back()); it != end(dir->children))
        {
            L_ERROR("Resource %s already exists, type is %d", path.c_str(), it->second->type);
            return false;
        } else {
            fs_file *node = new fs_file;
            node->type = fs_type::file;
            node->data_ptr = data;
            node->data_length = length;
            node->origin_length = origin_length;
            auto const ret = dir->children.emplace(std::pair{paths.back(), static_cast<fs_node_base*>(node)});
            return ret.second;
        }
    }
private:
    static std::vector<std::string> parse_path(const std::string &path) {
        // just split by character '/'
        return split_string(path, "/");
    }
};

static void free_node(fs_node_base *&node) {
    if (!node)
        return;
    auto const backup = node;
    node = nullptr; // 先置空
    if (backup->type == dir) {
        auto dir = static_cast<fs_dir*>(backup);
        // 先删除子项
        for (auto &child: dir->children) {
            free_node(child.second);
        }
        // 再删除自身
        delete dir;
    } else if (backup->type == file) {
        // 只需要简单的 delete
        delete static_cast<fs_file*>(backup);
    }
}

struct zstd_ctx {
    zstd_ctx() {
        ctx = ZSTD_createDCtx();
    }
    ~zstd_ctx() {
        ZSTD_freeDCtx(ctx);
    }

    ZSTD_DCtx *ctx;
};
thread_local zstd_ctx zstd;

static void decompress_data(const void *data_ptr, size_t data_length, size_t origin_length, std::vector<u8> &data_store)
{
    data_store.resize(origin_length);
    // TODO: using zstd to decompression
    auto const size = ZSTD_decompressDCtx(zstd.ctx, data_store.data(), origin_length, data_ptr, data_length);
    if (size != origin_length) {
        L_ERROR("ZSTD decompress failed!! %u != %u", origin_length, size);
        data_store.clear();
    }
}

embedded_resource::embedded_resource(void(*init_fn)(embedded_resource*)): d{new embedded_resource_data{}}
{
    if (init_fn)
        init_fn(this);
}

embedded_resource::~embedded_resource()
{
    fs_node_base *root = d->root;
    d->root = nullptr;
    free_node(root);
    delete d;
}

bool embedded_resource::register_resource(const std::string &path, const void *dat, size_t length, size_t origin_length)
{
    PI_ASSERT(path.starts_with('/'));
    PI_ASSERT(dat && length);
    // PI_ASSERT_MSG(origin_length >= length, "path: %s, length: 0x%zX, origin_length: 0x%zX", path.c_str(), length, origin_length); // BUG: 玄学压缩，这个不一定成立

    return d->set_file(path, dat, length, origin_length);
}

std::vector<std::string> embedded_resource::get_file_list(const std::string &dir)
{
    PI_ASSERT(dir.starts_with('/'));

    std::vector<std::string> result;

    auto const f_dir = d->get_dir(dir);
    if (f_dir) {
        std::transform(begin(f_dir->children), end(f_dir->children), std::back_inserter(result), [](auto const &it) {
            return it.first;
        });
    }
    return result;
}

std::vector<u8> embedded_resource::get_file_content(const std::string &path)
{
    PI_ASSERT(path.starts_with('/'));

    std::vector<u8> content;
    auto const f_file = d->get_file(path);
    if (f_file) {
        decompress_data(f_file->data_ptr, f_file->data_length, f_file->origin_length, content);
    }
    return content;
}
